import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cbook as cbook
import matplotlib.patches as mpatches
import csv
import random
import textwrap
import seaborn
import rawq
import pollutantdata as pd
import stormwaterutilities

def mapValues(inArray, scoresBasis, scores):
    outArray =  np.interp(inArray, rawq.tssScoreBasis, rawq.scores)
    return outArray

def getAnnotationPositions(x_data, y_data, txt_width, txt_height):
    a = zip(y_data, x_data)
    text_positions = y_data.copy()
    for index, (y, x) in enumerate(a):
        local_text_positions = [i for i in a if i[0] > (y - txt_height)
                            and (abs(i[1] - x) <= txt_width * 2)] #and i != (y,x)]
        if local_text_positions:
            sorted_ltp = sorted(local_text_positions)
            if abs(sorted_ltp[0][0] - y) < txt_height: #True == collision
                differ = np.diff(sorted_ltp, axis=0)
                a[index] = (sorted_ltp[-1][0] + txt_height, a[index][1])
                text_positions[index] = sorted_ltp[-1][0] + txt_height
                for k, (j, m) in enumerate(differ):
                    #j is the vertical distance between words
                    if j > txt_height * 2: #if True then room to fit a word in
                        a[index] = (sorted_ltp[k][0] + txt_height, a[index][1])
                        text_positions[index] = sorted_ltp[k][0] + txt_height
                        break
    return text_positions

def plotPollutant(pollutant, pollutantShortName, units, scoreBasis, concentration, reduction = None, outFile = None, colorBackground = False, annotate = True):
    #Populate magnitude score basis
    minScoreBasis = min(scoreBasis)
    maxScoreBasis = max(scoreBasis)
    fit_y = np.arange(minScoreBasis, maxScoreBasis, (maxScoreBasis-minScoreBasis)/100)
    fit_x = np.interp(fit_y, scoreBasis, pd.scores)

    #Populate landuse concentrations
    lst = zip(*concentration.values())
    lu = np.array(lst[0])
    mean = np.array(lst[1], dtype='f')
    lci = np.array(lst[2], dtype='f')
    uci = np.array(lst[3], dtype='f')

    mean_scores = np.interp(mean, scoreBasis, pd.scores)
    lci_scores = np.interp(lci, scoreBasis, pd.scores)
    uci_scores = np.interp(uci, scoreBasis, pd.scores)

    #Populate BMP efficacy
    if reduction is not None:
        lst = zip(*reduction)
        bmps = np.array(lst[0])
        bmp_mean = np.array(lst[1], dtype='f')
        bmp_low = np.array(lst[2], dtype='f')
        bmp_high = np.array(lst[3], dtype='f')

        bmp_mean_scores = np.interp(bmp_mean, scoreBasis, pd.scores)
        bmp_low_scores = np.interp(bmp_low, scoreBasis, pd.scores)
        bmp_high_scores = np.interp(bmp_high, scoreBasis, pd.scores)

    #Set some fancy styling
    seaborn.set_style('darkgrid', {'legend.frameon': True})
    seaborn.set_context('talk')
    seaborn.despine()
    luColor = 'dodgerblue'
    bmpColor = 'darkturquoise'
    basisColor = 'indigo'

    #Plot

    plt.close()

    plt.plot(fit_x, fit_y, '-', color=basisColor, ms=0, lw=1, label="Risk Score Basis")

    if reduction is not None and annotate == True:
        plt.errorbar(
            bmp_mean_scores, bmp_mean,
            yerr=[bmp_mean - bmp_low, bmp_high - bmp_mean],
            xerr=[bmp_mean_scores - bmp_low_scores, bmp_high_scores - bmp_mean_scores],
            label='BMP Effectiveness', fmt='D', color=bmpColor, ecolor=bmpColor, lw=1.0, ms=4)

    if annotate == True:
        plt.errorbar(
            mean_scores, mean,
            yerr=[mean - lci, uci - mean],
            xerr=[mean_scores - lci_scores, uci_scores - mean_scores],
            label='Land Use Concentrations', fmt='o', color=luColor, ecolor=luColor, lw=1.5, capsize=0, ms=6)

    if pollutant != pollutantShortName:
        plt.title("%s (%s) Pollutant Loading\nMagnitude of Failure Score Assignment"%(pollutant, pollutantShortName))
    else:
        plt.title("%s Pollutant Loading\nMagnitude of Failure Score Assignment"%(pollutant))

    #Add annotations
    #(labels, xs, ys) = zip(lu, mean_scores, mean)
    xaxisrange = plt.xlim()[1] - plt.xlim()[0]
    yaxisrange = plt.ylim()[1] - plt.ylim()[0]
    txt_width = 0.05*(xaxisrange)
    txt_height = 0.05*(yaxisrange)

    text_positions = getAnnotationPositions(mean_scores, mean,txt_width, txt_height)

    landuseNameDict = dict(zip(pd.landuse, pd.landuseLongName))

    if annotate == True:
        for (landuse, x, y, tp) in zip(lu, mean_scores, mean, text_positions):
            label = landuseNameDict[landuse]
            xlabelposition = np.clip(x+(1.5*xaxisrange*(abs(y-tp))/yaxisrange), plt.xlim()[0], plt.xlim()[1] - txt_width)
            #ylabelposition = np.clip(y-2*txt_height-2*txt_height*(abs(y-tp))/yaxisrange, plt.ylim()[0]+txt_height, plt.ylim()[1])
            ylabelposition = np.clip(tp, plt.ylim()[0]+txt_height, max(scoreBasis))
            plt.annotate(textwrap.fill(label, 10, break_long_words=False), xy=(x, y), xytext=(xlabelposition, ylabelposition), textcoords='data', ha='center', va='bottom',
                fontsize=13, bbox=dict(boxstyle='round,pad=0.2', fc=luColor, alpha=0.7),
                arrowprops=dict(arrowstyle='->', connectionstyle='angle,angleA=0,angleB=110,rad=10',  color='black')).draggable()

    txt_height = 0.05*(plt.ylim()[1] - plt.ylim()[0])
    txt_width = 0.15*(plt.xlim()[1] - plt.xlim()[0])

    if reduction is not None and annotate == True:
        text_positions = getAnnotationPositions(bmp_mean_scores, bmp_mean, txt_width, txt_height)
        for (label, x, y, tp) in zip(bmps, bmp_mean_scores, bmp_mean, text_positions):
            xlabelposition = x #np.clip(x-txt_width, plt.xlim()[0], plt.xlim()[1] - txt_width)
            ylabelposition = np.clip(tp, plt.ylim()[0]+txt_height, max(scoreBasis))

            plt.annotate(textwrap.fill(label, 10, break_long_words=False), xy=(x, y), xytext=(xlabelposition, ylabelposition), textcoords='data', ha='center', va='bottom',
                fontsize=12, bbox=dict(boxstyle='round,pad=0.2', fc=bmpColor, alpha=0.7, zorder=10),
                arrowprops=dict(arrowstyle='->', connectionstyle='angle,angleA=0,angleB=110,rad=10',  color='black')).draggable()

    #Labeling
    plt.xlabel('Magnitude of Failure Score')
    plt.ylabel('%s (%s)'%(pollutantShortName,units))

    if reduction is not None:
        ax = plt.axis([1, 5.01, 0, 1.1 * max(max(uci), max(scoreBasis))])
    else:
        ax = plt.axis([1, 5.01, 0, max(max(uci), max(scoreBasis))])


    handles, labels = plt.gca().get_legend_handles_labels()
    display = (0,1,2)

    lowRiskPatch = mpatches.Patch(color='mediumseagreen', label='Concentration Similar to\nPristine')
    mediumRiskPatch = mpatches.Patch(color='khaki', label='Water Quality\nStandard')
    highRiskPatch = mpatches.Patch(color='salmon', label='Acute\nToxicity')

    plt.legend([handle for i,handle in enumerate(handles) if i in display]+[lowRiskPatch,mediumRiskPatch,highRiskPatch],
          [label for i,label in enumerate(labels) if i in display]+['Similar to Pristine\nConditions', 'Water Quality\nStandard', 'Acute\nToxicity'],
          loc='upper left', shadow=True, fancybox=True)
    #legend2 = plt.legend([red_patch])



    #Color the background
    if colorBackground:
        from matplotlib import colors

        x = np.arange(1, 6, .1)

        y = np.array(plt.yticks()[0])
        (X, Y) = np.meshgrid(x,y)
        C = np.array([x for i in range(len(y))])
        cm = colors.LinearSegmentedColormap.from_list('Risk',
            ['mediumseagreen', 'khaki', 'khaki', 'sandybrown', 'salmon'])

        xmin, xmax = (plt.xlim()[0], plt.xlim()[1])
        ymin, ymax = (plt.ylim()[0], plt.ylim()[1])

        xbox = np.append(fit_x, [5.1, 5.1, 0])
        ybox = np.append(fit_y, [max(fit_y), 0, 0])
        poly, = plt.fill(xbox, ybox, facecolor='None', closed=True)
        img_data = np.array(np.arange(ymin,ymax,(ymax-ymin)/100.), ndmin=2)

        #img_data = img_data.reshape(img_data.size,1)

        im = plt.imshow(img_data, aspect='auto', origin='lower', cmap=cm, extent=[xmin,xmax,ymin,ymax])#, vmin=0, vmax=fit_y.max())

        im.set_clip_path(poly)



        plt.grid(b=True, which='major', axis='both', linewidth=.5, color='dimgray', zorder=1)

    if outFile == None:
        plt.show()
    else:
        plt.savefig(outFile)


applyAnnotation = False


plotPollutant('Total Suspended Sediment', 'TSS', 'mg/l', pd.tssScoreBasis, pd.tssLuConcentration, pd.tssReduction, colorBackground = True, annotate = applyAnnotation)
plotPollutant('Dissolved Lead', 'PbD', 'ug/l', pd.pbScoreBasis, pd.pbLuConcentration, pd.pbReduction, colorBackground = True, annotate = applyAnnotation)
plotPollutant('Total Phosphorus', 'TP', 'mg/l', pd.tpScoreBasis, pd.tpLuConcentration, pd.tpReduction, colorBackground = True, annotate = applyAnnotation)
plotPollutant('Biochemical Oxygen Demand', 'BOD', 'mg/l', pd.pbScoreBasis, pd.pbLuConcentration, pd.pbReduction, colorBackground = True, annotate = applyAnnotation)
plotPollutant('E. Coli', 'E. Coli', 'CFU / 100 ml', pd.ecoliScoreBasis, pd.ecoliLuConcentration, pd.ecoliReduction, colorBackground = True, annotate = applyAnnotation)


#plotPollutant('Total Suspended Sediment', 'TSS', 'mg/l', rawq.tssScoreBasis, pd.tssLuConcentration, pd.tssReduction, colorBackground = True, annotate = False)
#plotPollutant('Dissolved Lead', 'PbD', 'ug/l', rawq.pbScoreBasis, pd.pbLuConcentration, pd.pbReduction, colorBackground = True, annotate = False)
#plotPollutant('Total Phosphorus', 'TP', 'mg/l', rawq.tpScoreBasis, pd.tpLuConcentration, pd.tpReduction, colorBackground = True, annotate = False)
#plotPollutant('Biochemical Oxygen Demand', 'BOD', 'mg/l', rawq.pbScoreBasis, pd.pbLuConcentration, pd.pbReduction, colorBackground = True, annotate = False)
#plotPollutant('E. Coli', 'E. Coli', 'CFU / 100 ml', rawq.ecoliScoreBasis, pd.ecoliLuConcentration, pd.ecoliReduction, colorBackground = True, annotate = False)
